

import asyncio
import logging
import op_coroutine  as co
from calculate_main_star_collector import Collector
from calculate_main_star_dispatcher import Dispatcher

import keti4_config as config


async def dispatch_task(file="data.png", target_nodes=[0x02, 0x03, 0x04, 0x05], send_port=8800):
    # 2. 计算
    dispatcher = Dispatcher(send_port=send_port, target_nodes=target_nodes, file=file)
    await dispatcher.run()
    return True


def collect_data(recv_port=8800):
    # 1. 接收计算主星发送的数据
    collector = Collector(recv_port=recv_port, recv_num=4)
    data = collector.run()
    return data


async def main(file="data.png", target_nodes=[0x02, 0x03, 0x04, 0x05], send_port=16536, recv_port=8800):
    tasks = [
        # (collect_data, [], {'recv_port':recv_port}),
        # (dispatch_task, [], {'file':file, 'target_nodes':target_nodes, 'recv_port':recv_port})
        (collect_data, [recv_port], {}),
        (dispatch_task, [file, target_nodes, send_port], {})
    ]
    result = await co.go_all(tasks)
    print("All Results:", result)


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)

    # 获取命令行参数
    # script_name = sys.argv[0]  # 脚本名
    # arguments = sys.argv[1:]   # 除脚本名外的参数列表

    # data = read_data(sys.argv[1])   # 除脚本名外的参数列表
    # target_node = sys.argv[2]  # 脚本名

    filename = "data.png"
    # recv_port = 8800  # 目标节点名称

    target_nodes = [0x02, 0x03, 0x04, 0x05]  # 目标节点名称

    try:
        asyncio.run(main(file=filename, 
                         target_nodes=target_nodes, 
                         send_port=config.send_port ,
                         recv_port=config.recv_port))
    except Exception as e:
        print("Error:", e)
        raise e
    finally:
        co.shutdown()
        print("Done")
